# 一方库
from myutils.helper import get_mean


# 官方库
import glob
import glob
import os
import datetime

# 三方库
import torch



# 模型保存，文件夹/model_twin_时间_精度.pth
# time_curr_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
# msg='out'
# path_model = './model_one_{}_{}_{:02d}.pth'.format(time_curr_str,msg,int(list_accuracy_Y[-1]*100))
# torch.save(model.state_dict(), path_model)


def get_path_model_dir(model_name):
    model_dir = f'checkpoint/{model_name}'
    return model_dir


def model_load(model,model_name=None,path_model=None,device='cpu'):
    """模型载入（按照模型名或模型路径） 
    
    Args:
        model_name: 模型名
        path_model_s：所有和模型名匹配的路径
        path_model：模型路径（若不存在，则为 path_model_s 中最新的）
    """
    
    if path_model is None:
        try:

            path_model_s = glob.glob(f'{get_path_model_dir(model_name)}/{model_name}*.pth')
            path_model = max(path_model_s)
        except:
            print(f'不存在已训练的模型:{model_name}')
            return
    
    model_param = torch.load(path_model,map_location=device)
    result = model.load_state_dict(model_param)
    print(f'model_load:{path_model}')
    print(result)

def model_save(model,model_name,print_margin=100,list_accuracy_Y=[]):
    """模型保存（按照模型名）
    
    model: 模型
    model_name: 模型名
    print_margin: 取 精度数组的 最后n个
    """
    time_curr_str = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    accuracy = get_mean(list_accuracy_Y,print_margin) if len(list_accuracy_Y)!=0 else 0

    
    path_model = f'{get_path_model_dir(model_name)}'+'/{}_{}_{:03d}.pth'.format(model_name,time_curr_str,int(accuracy*1000))
    path_model_dir = os.path.dirname(path_model)

    if not os.path.exists(path_model_dir):
        os.mkdir(path_model_dir)
    
    torch.save(model.state_dict(), path_model)
    print(f'model_save:{path_model}')